import argparse
import logging
import os

from deap import tools

from plotting.drawer import plot_gantt_chart, draw_precedence_relations
from solution_methods.helper_functions import load_parameters, load_job_shop_env
from solution_methods.GA.src.operators import (evaluate_individual, evaluate_population, repair_precedence_constraints, variation)
from solution_methods.GA.utils import record_stats, output_dir_exp_name, results_saving
from solution_methods.GA.src.initialization import initialize_run

logging.basicConfig(level=logging.INFO)
PARAM_FILE = "../../configs/GA.toml"


def run_GA(jobShopEnv, population, toolbox, stats, hof, **kwargs):
    """
    执行遗传算法并返回最佳个体。

    参数:
        jobShopEnv: 问题环境。
        population: 初始种群。
        toolbox: DEAP工具箱。
        stats: DEAP统计信息。
        hof: Hall of Fame（名人堂）。
        kwargs: 额外的关键字参数。

    返回:
        tuple: 最佳个体的makespan和作业车间环境。
    """

    # 初始化名人堂和统计信息
    hof.update(population)

    gen = 0
    logbook = tools.Logbook()
    logbook.header = ["gen"] + (stats.fields if stats else [])
    df_list = []

    # 记录初始种群的统计信息
    record_stats(gen, population, logbook, stats, kwargs['output']['logbook'], df_list, logging)
    if kwargs['output']['logbook']:
        logging.info(logbook.stream)

    for gen in range(1, kwargs['algorithm']['ngen'] + 1):
        # 对种群进行变异操作
        offspring = variation(population, toolbox,
                              pop_size=kwargs['algorithm'].get('population_size'),
                              cr=kwargs['algorithm'].get('cr'),
                              indpb=kwargs['algorithm'].get('indpb'))

        # 如果环境需要，修复优先约束关系（仅适用于装配调度）
        if any(keyword in jobShopEnv.instance_name for keyword in ['/dafjs/', '/yfjs/']):
            try:
                offspring = repair_precedence_constraints(jobShopEnv, offspring)
            except Exception as e:
                logging.error(f"修复优先约束关系时出错: {e}")
                continue

        # 评估后代的适应度
        try:
            fitnesses = evaluate_population(toolbox, offspring)
            for ind, fit in zip(offspring, fitnesses):
                ind.fitness.values = fit
        except Exception as e:
            logging.error(f"评估后代适应度时出错: {e}")
            continue

        # 选择下一代种群
        population[:] = toolbox.select(population + offspring)

        # 更新名人堂和统计信息
        hof.update(population)
        record_stats(gen, population, logbook, stats, kwargs['output']['logbook'], df_list, logging)

    # 评估最佳个体的makespan
    makespan, jobShopEnv = evaluate_individual(hof[0], jobShopEnv, reset=False)
    logging.info(f"Makespan: {makespan}")
    return makespan, jobShopEnv


def main(param_file=PARAM_FILE):
    """
    主函数，加载参数文件，初始化作业车间环境和遗传算法，并运行遗传算法。

    参数:
        param_file: 参数文件路径，默认为PARAM_FILE。
    """
    try:
        parameters = load_parameters(param_file)
        logging.info(f"参数已从 {param_file} 加载。")
    except FileNotFoundError:
        logging.error(f"未找到参数文件 {param_file}。")
        return

    # 加载作业车间环境，并初始化遗传算法
    jobShopEnv = load_job_shop_env(parameters['instance'].get('problem_instance'))
    population, toolbox, stats, hof = initialize_run(jobShopEnv, **parameters)
    makespan, jobShopEnv = run_GA(jobShopEnv, population, toolbox, stats, hof, **parameters)

    if makespan is not None:
        # 检查输出配置并准备输出路径（如果需要）
        output_config = parameters['output']
        save_gantt = output_config.get('save_gantt')
        save_results = output_config.get('save_results')
        show_gantt = output_config.get('show_gantt')
        show_precedences = output_config.get('show_precedences')

        if save_gantt or save_results:
            output_dir, exp_name = output_dir_exp_name(parameters)
            output_dir = os.path.join(output_dir, f"{exp_name}")
            os.makedirs(output_dir, exist_ok=True)

        # 绘制优先关系图（如果需要）
        if show_precedences:
            draw_precedence_relations(jobShopEnv)

        # 绘制甘特图（如果需要）
        if show_gantt or save_gantt:
            logging.info("正在生成甘特图。")
            plt = plot_gantt_chart(jobShopEnv)

            if save_gantt:
                plt.savefig(output_dir + "/gantt.png")
                logging.info(f"甘特图已保存到 {output_dir}")

            if show_gantt:
                plt.show()

        # 保存结果（如果启用）
        if save_results:
            results_saving(makespan, output_dir, parameters)
            logging.info(f"结果已保存到 {output_dir}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="运行遗传算法")
    parser.add_argument(
        "-f", "--config_file",
        type=str,
        default=PARAM_FILE,
        help="配置文件路径",
    )
    args = parser.parse_args()
    main(param_file=args.config_file)
